import argparse

class ArgumentParser:
    def __init__(self):
        self.parser = argparse.ArgumentParser(description='Mini AI Translator.')
        self.parser.add_argument('--config', type=str, default='config.yaml',
                                 help='Configuration file with model and API settings.')
        self.parser.add_argument('--supported_models', type=str, default='models.yaml',
                                 help='File containing a list of supported models.')
        self.parser.add_argument('--model_type', type=str, choices=['OpenAI', 'ZhipuAI', 'Tongyi'],
                                 help='Type of the translation model to use (e.g., OpenAI, ZhipuAI, Tongyi).')
        self.parser.add_argument('--model', type=str,
                                 help='The specific model to use for translation. For example, "gpt-3.5-turbo" for OpenAI or "glm-4" for ZhipuAI or "qwen-turbo" for Tongyi.')
        self.parser.add_argument('--target_language', type=str, default='Chinese',
                                 help='The target language for translation.')
        self.parser.add_argument('--input_file_path', type=str,
                                 help='PDF file to translate.')
        self.parser.add_argument('--output_file_path', type=str, default='./translated_books',
                                 help='Path to save the translated book.')
        self.parser.add_argument('--output_file_format', type=str, choices=['markdown'], default='markdown',
                                 help='The file format of the translated book. Currently, only Markdown format is supported.')

    def parse_arguments(self):
        return self.parser.parse_args()


# Test argument_parser
if __name__ == '__main__':
    # python3 argument_parser.py --model_type OpenAI --model gpt-3.5-turbo --input_file_path ../jupyter/test.pdf --output_file_path ./translated_books --output_file_format markdown
    # python3 argument_parser.py --model_type ZhipuAI --model glm-4 --input_file_path ../jupyter/test.pdf --output_file_path ./translated_book
    # python3 argument_parser.py --model_type Tongyi --model qwen-turbo --input_file_path ../jupyter/test.pdf --output_file_path ./translated_book
    parser = ArgumentParser()
    args = parser.parse_arguments()
    print(f"args = {args}")